// kdtree_build.cpp
// 04/17/2011: created

#include <fstream>
#include <cassert>
#include <limits>
#include <algorithm>
#include <queue>
#include <QtGui>
#include "geometric_problem.h"
#include "Log.h"
#include "kdtree.h"
#include "covertree.h"

namespace common {
  void loadVVec(const std::string& fName, std::vector<vec>& vvec) {
    int n, d;
    std::ifstream is(fName.c_str());
    is >> n >> d;

    Log::m_log << "n=" << n << " d=" << d << endl;
    vvec.clear();
    for (int i = 0; i < n && !is.eof(); i++) {
      vvec.push_back(vec(d));
      vec& v = vvec.back();
      v.load(is, arma::arma_ascii);
      assert(v.n_elem == d);
      //Log::m_log << "v[" << i << "]" << endl;
    }
    is.close();
  }
  void drawVVec(const std::vector<vec>& vvec, Plotter& pl, const QPen& pen) {
    pl.setPen(pen);
    int n = vvec.size();
    for (int i = 0; i < n; i++) {
      const vec& p = vvec[i];
      pl.drawEllipse(pl.point(p), 2.0, 2.0);
    }
  }
}

void KDTreeBuild::loadData(const std::string& fName)
{
  if (m_tree) { delete m_tree; m_tree = NULL; }
  common::loadVVec(fName, m_point);
}

void KDTreeBuild::solve()
{
  int n = m_point.size();
  if (m_tree) { delete m_tree; m_tree = NULL; }
  m_tree = new KDTree(m_point);
  m_tree->buildTree();
}

struct BoundingBoxDrawer 
{
  Plotter& m_pl;
  BoundingBoxDrawer(Plotter& pl) : m_pl(pl) {}
  void operator()(Tree* tree_node) {
    KDTree *node = static_cast<KDTree*>(tree_node);
    const BoundingBox& b = node->box();
    m_pl.drawRect(QRectF(m_pl.point(b.min()), m_pl.point(b.max())));
  }
};

void KDTreeBuild::draw(Plotter& pl)
{
  QPen pen(Qt::SolidLine);
  pen.setWidth(2);

  pen.setColor(Qt::black);
  common::drawVVec(m_point, pl, pen);

  if (m_tree) {
    pen.setColor(Qt::red);
    pen.setStyle(Qt::DashLine);
    pen.setWidth(1);
    pl.setPen(pen);
    
    BoundingBoxDrawer drawer(pl);
    m_tree->visitRecursive(drawer);
  }
}

void CoverTreeBuild::loadData(const std::string& fName)
{
  if (m_tree) { delete m_tree; m_tree = NULL; }
  common::loadVVec(fName, m_point);
}

void CoverTreeBuild::solve()
{
  int n = m_point.size();
  if (m_tree) { delete m_tree; m_tree = NULL; }
  m_tree = new CoverTree(m_point);
  m_tree->buildTree();
}

struct BallDrawer 
{
  Plotter& m_pl;
  BallDrawer(Plotter& pl) : m_pl(pl) {}
  void operator()(Tree* tree_node) {
    CoverTree *node = static_cast<CoverTree*>(tree_node);
    double r = pow(2, node->log_radius());
    QPointF R = m_pl.distance(r, r);
    //Log::m_log << "r.x=" << R.x() << " r.y=" << R.y() << endl;
    m_pl.drawEllipse(m_pl.point(node->get(0)), R.x(), -R.y());
  }
};

void CoverTreeBuild::draw(Plotter& pl)
{
  QPen pen(Qt::SolidLine);
  pen.setWidth(2);

  pen.setColor(Qt::black);
  common::drawVVec(m_point, pl, pen);

  if (m_tree) {
    pen.setColor(Qt::red);
    pen.setStyle(Qt::DotLine);
    pen.setWidth(1);
    pl.setPen(pen);

    BallDrawer drawer(pl);
    m_tree->visitRecursive(drawer);
  }
}

void Matching::loadData(const std::string& fName1, const std::string& fName2)
{
  common::loadVVec(fName1, m_point1);
  common::loadVVec(fName2, m_point2);
  m_match.clear();
}

void Matching::draw(Plotter& pl)
{
  QPen pen(Qt::SolidLine);

  pen.setWidth(1);
  pen.setColor(Qt::black);
  pen.setStyle(Qt::DashLine);
  pl.setPen(pen);
  int n = n1();
  for (int x = 0; x < n && x < (int) m_match.size(); x++) {
    int y = m_match[x]-n;
    if (y >= 0) {
      pl.drawLine(pl.point(m_point1[x]), pl.point(m_point2[y]));
    }
  }

  pen.setWidth(2);
  pen.setStyle(Qt::SolidLine);

  pen.setColor(Qt::blue);
  common::drawVVec(m_point1, pl, pen);

  pen.setColor(Qt::red);
  common::drawVVec(m_point2, pl, pen);

}

void Matching::solve()
{
  assert(n1() == n2());
  n = n1();
  m_match = std::vector<int>(2*n, -1);// init M = \emptyset
  m_p = std::vector<double>(2*n, 0);  // init p(x) = p(y) = 0, for all x \in X, y \in Y
                                      // this satisfies c_p(x,y) = p(x)+c(x,y)-p(y) >= 0, forall x,y
  for (int matching_size = 0; matching_size < n; ) {
    int t = compute_shortest_distances();     // compute shortest distances from s to all v in G_M (c(x,y)+p(x)-p(y))
    //Log::m_log << "free y=" << t << endl;
    augment_path(t);
    for (unsigned int i = 0; i < m_p.size(); i++) {
      m_p[i] += m_dist[i];
      //Log::m_log << "p[" << i << "]=" << m_p[i] << endl;
    }
    matching_size++;
    if (matching_size%100==0) Log::m_log << "size=" << matching_size << endl;
  }
  Log::m_log << "cost=" << cost() << endl;
}

// Dijkstra algorithm, applied on G_M c_p(x,y) = c(x,y)+p(x)-p(y) >= 0
int Matching::compute_shortest_distances()
{
  m_prev = std::vector<int>(2*n, -1);
  m_dist = std::vector<double>(2*n, POSITIVE_INFINITY);
  std::vector<bool> visited(2*n, false);
  for (int x = 0; x < n; x++) if (m_match[x] == -1)// all unassigned left nodes
    m_dist[x] = 0;                                 // has zero distance from s
  int t = -1;
  while (1) {                                      // repeat until all distances are computed
    double min_dist = POSITIVE_INFINITY;
    int min_u = -1;
    for (int i = 0; i < 2*n; i++) {
      //Log::m_log << i << "=(" << visited[i] << "," << m_dist[i] << ") ";
      if (!visited[i] && m_dist[i] < min_dist) {
        min_u = i;
        min_dist = m_dist[i];
      }
    }
    //Log::m_log << endl << "min_u = " << min_u << " min_dist=" << min_dist << endl;
    if (min_u == -1) break;                        // done

    visited[min_u] = true;
    if (min_u >= n) {                              // a right node, using matching edge
      int y = min_u, x = m_match[y];
      if (x == -1) {                               // a free y \in Y, shortest path from s to t found
        if (t == -1) t = y;
      }
      else {
        double c_p = weight_p(x,y), new_dist = m_dist[y] + c_p;
        //Log::m_log << "(y,x)=(" << y << "," << x << ") p[x]=" << m_p[x] << " w=" << weight(x,y-n) << " p[y]=" << m_p[y] << endl;
        if (m_dist[x] > new_dist) {
          m_prev[x] = y;
          m_dist[x] = new_dist;
          //Log::m_log << "update (y,x)=(" << y << "," << x << ") cp=" << c_p << " d=" << new_dist << endl;
        }
      }
    }
    else {                                         // a left node
      int x = min_u;
      for (int y = n; y < 2*n; y++) {              // for all right nodes
        double c_p = weight_p(x,y), new_dist = m_dist[x] + c_p;
        //Log::m_log << "(x,y)=(" << x << "," << y << ") p[x]=" << m_p[x] << " w=" << weight(x,y-n) << " p[y]=" << m_p[y] << endl;
        if (!visited[y] && m_dist[y] > new_dist) {
          m_prev[y] = x;
          m_dist[y] = new_dist;
          //Log::m_log << "update (x,y)=(" << x << "," << y << ") cp=" << c_p << " d=" << new_dist << endl;
        }
      }
    }
  }
  return t;
}

void Matching::augment_path(int t)
{
  //Log::m_log << "Augmenting path: " << t;
  while (t != -1) {
    int x = m_prev[t];
    m_match[x] = t;
    m_match[t] = x;
    t = m_prev[x];
    //Log::m_log << " --> " << x << " --> " << t;
  }
  //Log::m_log << endl;
}

double Matching::cost() const
{
  //Log::m_log << "compute cost " << endl;
  int n = n1();
  double s = 0;

  // sanity check
  std::vector<bool> used(n, false);
  for (int x = 0; x < n; x++) {
    int y = m_match[x]-n;
    if (used[y]) {
      Log::m_log << "bad solution" << endl;
    }
    used[y] = true;
  }

  for (int x = 0; x < n && x < (int) m_match.size(); x++) {
    int y = m_match[x]-n;
    if (y >= 0) {
      s += weight(x, y);
    }
  }
  return s;
}

void AuctionMatching::solve()
{
  assert(n1() == n2());
  n = n1();
  m_p = std::vector<double>(2*n, 0);  // init p(x) = p(y) = 0, for all x \in X, y \in Y
                                      // this satisfies (x,y)\in M => w(x,y)+p(y) <= min w(x,y) + p(y) + epsilon
  std::queue<int> bidders;
  for (double epsilon = n/10.0; epsilon >= 1.0/n || epsilon >= 1e-3; epsilon /= 5) {
    m_match = std::vector<int>(2*n, -1);  // init M = \emptyset, keep prices from previous iteration
    for (int i = 0; i < n; i++) bidders.push(i);
    int match_size=0;
    Log::m_log << "---- epsilon = " << epsilon << " ----" << endl;
    while (!bidders.empty()) {
      int x = bidders.front(), min_y = -1; bidders.pop();

      double v = POSITIVE_INFINITY, w = POSITIVE_INFINITY;
      for (int y = n; y < 2*n; y++) {
        double cost = weight(x, y-n)+m_p[y];
        if (cost < v) { w = v; v = cost; min_y = y; }
        else if (cost < w) { w = cost; }
      }
      m_p[min_y] += w - v + epsilon;

      if (m_match[min_y] != -1) {
        bidders.push(m_match[min_y]);
        m_match[m_match[min_y]] = -1;
        //Log::m_log << "push " << m_match[min_y] << " "; Log::m_log.flush();
      }
      else {
        match_size++;
        if (match_size%100 == 0) Log::m_log << "size=" << match_size << endl;
      }
      m_match[x] = min_y;
      m_match[min_y] = x;
    }
  }
  Log::m_log << endl << "cost=" << cost() << endl;
}
